import copy
from math import floor, sqrt

import torch

from algorithms.algorithm import Algorithm
from algorithms.base.nsga2 import NSGA2
from operators.mutations.polynomial_mutation import PolynomialMutation
from problems.problem import Problem
from utils.sort import sortrows
from utils.vectors import get_ref_vectors, vector_radian


class DGEA(Algorithm):
    parameters = [
        {'label': 'RefNo', 'name': 'ref_no', 'type': 'number',
         'description': 'Number of reference vectors for offspring generation', 'step': 1,
         'defaultValue': 10},
        {
            'label': 'operation', 'name': 'operation', 'type': 'select',
            'description': 'Operation of the environmental selection', 'defaultValue': 0,
            "options": [
                {'label': "subRVEA", "value": 0},
                {'label': "subNSGAII", "value": 1},
                {'label': "subIBEA", "value": 2},
                {'label': "subSPEA2", "value": 3}]
        }
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs,
                 ):
        self.v = get_ref_vectors(pop_size, problem.obj_dim)
        pop_size = len(self.v)
        self.population = problem.init_population(pop_size)
        self.offspring = problem.init_population(pop_size)
        self.arc = copy.deepcopy(self.population)
        super().__init__(pop_size, problem, kwargs)

    def each_iteration(self):
        self.population, front_no = self.pre_selection(self.population + self.offspring, self.v,
                                                       (self.problem.fun_eval / self.problem.max_fun_eval) ** 2,
                                                       self.ref_no)
        self.offspring = self.direction_reproduction(self.problem, self.population, front_no, self.ref_no,
                                                     self.pop_size)

        match self.operation:
            case 0:
                self.arc = self.sub_RVEA(self.arc + self.offspring, self.v,
                                         (self.problem.fun_eval / self.problem.max_fun_eval) ** 2)
            case 1:
                self.arc = self.sub_NSGAII(self.arc + self.offspring, self.pop_size)
            case 2:
                self.arc = self.sub_IBEA(self.arc + self.offspring, self.pop_size, 0.05)
            case 3:
                self.arc = self.sub_SPEA2(self.arc + self.offspring, self.pop_size)
            case _:
                self.arc = copy.deepcopy(self.population) + copy.deepcopy(self.offspring)

    def final_output(self):
        return self.arc

    @staticmethod
    def sub_RVEA(population, reference_vectors, theta):
        pop_obj = population.obj
        N, M = pop_obj.shape
        NV = reference_vectors.shape[0]

        # Translate the population
        min_values = torch.min(pop_obj, dim=0)[0]
        pop_obj = pop_obj - min_values

        # Calculate the degree of violation of each solution
        CV = torch.sum(torch.clamp(population.con, min=0), dim=1)

        # Calculate the smallest angle value between each vector and others
        gamma = vector_radian(reference_vectors, reference_vectors,
                              return_radian=True).fill_diagonal_(torch.inf).min(1)[0]

        # Associate each solution to a reference vector
        angles = vector_radian(pop_obj, reference_vectors, return_radian=True)
        associate = torch.min(angles, dim=1)[1]

        # Select one solution for each reference vector
        Next = torch.full((NV,), -1, dtype=torch.long)
        unique_associates = torch.unique(associate)

        for i in unique_associates:
            current1 = torch.where((associate == i) & (CV == 0))[0]
            current2 = torch.where((associate == i) & (CV != 0))[0]

            if current1.numel() > 0:
                # Calculate the APD value of each solution
                APD = (1 + M * theta * angles[current1, i] / gamma[i]) * torch.sqrt(
                    torch.sum(pop_obj[current1] ** 2, dim=1))
                # Select the one with the minimum APD value
                best = torch.argmin(APD)
                Next[i] = current1[best]
            elif current2.numel() > 0:
                # Select the one with the minimum CV value
                best = torch.argmin(CV[current2])
                Next[i] = current2[best]

        # Population for next generation
        return population[Next[Next != -1]]

    @staticmethod
    def sub_NSGAII(population, pop_size):
        return population[NSGA2.environmental_selection(population.obj, pop_size)]

    @classmethod
    def sub_IBEA(cls, population, N, kappa=0.05):
        next_indices = torch.arange(len(population))
        fitness, I_matrix, C_values = cls.calculate_fitness0(population.obj, kappa)

        while len(next_indices) > N:
            min_fitness_index = torch.argmin(fitness[next_indices])
            x = next_indices[min_fitness_index]
            fitness = fitness + torch.exp(-I_matrix[x] / C_values[x] / kappa)
            next_indices = next_indices[next_indices != x]

        return population[next_indices]

    @staticmethod
    def calculate_fitness0(pop_objs, kappa):
        """
        Calculate the fitness values for IBEA (in Python).

        Parameters:
        - pop_objs (Tensor): The objective values of the population.
        - kappa (float): The kappa parameter.

        Returns:
        - fitness (Tensor): The fitness values of the population.
        - I_matrix (Tensor): The `I` matrix used in IBEA.
        - C_values (Tensor): The `C` values used in IBEA.
        """
        N = pop_objs.shape[0]

        # Normalize objective values
        min_objs = torch.min(pop_objs, dim=0)[0]
        max_objs = torch.max(pop_objs, dim=0)[0]
        normal_objs = (pop_objs - min_objs) / (max_objs - min_objs)

        # Calculate I matrix
        I_matrix = torch.zeros((N, N))
        for i in range(N):
            for j in range(N):
                I_matrix[i, j] = torch.max(normal_objs[i] - normal_objs[j]).item()

        # Calculate C values
        C_values = torch.max(torch.abs(I_matrix), dim=0)[0]

        # Calculate fitness values
        fitness = torch.sum(-torch.exp(-I_matrix / C_values / kappa), dim=0) + 1

        return fitness, I_matrix, C_values

    @classmethod
    def sub_SPEA2(cls, population, N):
        """
            The environmental selection of SPEA2 (in Python).

            Parameters:
            - population (list): The population objects.
            - N (int): The number of solutions to select.

            Returns:
            - selected_population (list): The selected population for the next generation.
            """
        # Calculate the fitness of each solution
        fitness = cls.calculate_fitness(population.obj)
        CC = min(len(population), N)

        # Environmental selection
        next_selected = fitness < 1
        if torch.sum(next_selected) < CC:
            rank = torch.argsort(fitness)
            next_selected[rank[:CC]] = True
        elif torch.sum(next_selected) > CC:
            del_indices = cls.truncation(population.obj[next_selected], torch.sum(next_selected) - CC)
            temp = torch.where(next_selected)[0]
            next_selected[temp[del_indices]] = False
        # Population for next generation
        return population[next_selected]

    @staticmethod
    def truncation(pop_objs, K):
        """
        Select part of the solutions by truncation.

        Parameters:
        - pop_objs (Tensor): The objective values of the population.
        - K (int): The number of solutions to truncate.

        Returns:
        - del_indices (list): The indices of solutions to be truncated.
        """
        # Truncation
        distance = torch.cdist(pop_objs, pop_objs)
        distance.fill_diagonal_(torch.inf)
        del_mask = torch.zeros(pop_objs.shape[0], dtype=bool)

        while del_mask.sum() < K:
            remain_indices = torch.where(~del_mask)[0]
            temp = distance[remain_indices]
            sorted_distance = torch.sort(temp, dim=1)[0]
            rank = sortrows(sorted_distance, return_sorted_indices=True)[1]
            del_mask[remain_indices[rank[0]]] = True

        return del_mask

    @staticmethod
    def calculate_fitness(pop_objs):
        """
        Calculate the fitness of each solution.

        Parameters:
        - pop_objs (Tensor): The objective values of the population.

        Returns:
        - fitness (Tensor): The fitness values of the population.
        """
        N = pop_objs.shape[0]

        # Detect the dominance relation between each two solutions
        dominate = torch.zeros((N, N), dtype=bool)
        for i in range(N - 1):
            for j in range(i + 1, N):
                k = torch.any(pop_objs[i] < pop_objs[j]).to(torch.int) - torch.any(pop_objs[i] > pop_objs[j]).to(
                    torch.int)
                if k == 1:
                    dominate[i, j] = True
                elif k == -1:
                    dominate[j, i] = True

        # Calculate S(i)
        S = dominate.sum(dim=1)

        # Calculate R(i)
        R = torch.zeros(N)
        for i in range(N):
            R[i] = S[dominate[:, i]].sum()

        # Calculate D(i)
        distance = torch.cdist(pop_objs, pop_objs)
        distance.fill_diagonal_(torch.inf)
        sorted_distance = torch.sort(distance, dim=1)[0]
        D = 1.0 / (sorted_distance[:, floor(sqrt(N))] + 2.0)

        # Calculate the fitnesses
        fitness = R + D
        return fitness

    @staticmethod
    def pre_selection(population, v, theta, ref_no):
        NV = len(v)
        obj = population.obj
        front_no, max_front = NSGA2.ND_sort(population.obj, min(NV, len(population)))
        if torch.sum(front_no == 1) >= ref_no:
            next0 = front_no < max_front
            last = torch.where(front_no == max_front)[0]
        else:
            next0 = front_no == 1
            last = torch.where(front_no > 1)[0]
        obj = obj[last]
        N, M = obj.shape
        obj = obj - obj.min(0)[0]
        gamma = vector_radian(v, v, return_radian=True)
        gamma.fill_diagonal_(torch.inf)
        gamma = gamma.min(1)[0]

        angle = vector_radian(obj, v, return_radian=True)
        associate = angle.min(1)[1]

        # Select one solution for each reference vector
        next_index = torch.full((NV,), -1, dtype=torch.long)
        for i in torch.unique(associate):
            current = torch.where(associate == i)[0]
            APD = (1 + M * theta * angle[current, i] / gamma[i]) * torch.sqrt(torch.sum(obj[current] ** 2, dim=1))
            best = torch.argmin(APD)
            next_index[i] = current[best]
        next0[last[next_index[next_index != -1]]] = True
        return population[next0], front_no[next0]

    @staticmethod
    def direction_reproduction(problem, population, front_no, ref_no, pop_size):
        dec = population.dec
        N, D = dec.shape
        indexD = torch.where(front_no > 1)[0]
        domiN = torch.numel(indexD)
        indexN = front_no == 1
        non_dec = dec[indexN]
        domi_dec = dec[indexD]
        sub_N = max(floor(pop_size / ref_no), 10)
        startP = non_dec[torch.randperm(N - domiN)][0]
        if domiN < ref_no:
            if N <= ref_no:
                endP = torch.cat((domi_dec, non_dec[torch.randperm(N - domiN)]), dim=0)
            else:
                endP = torch.cat((domi_dec, non_dec[torch.randperm(N - domiN)[:ref_no - domiN]]), dim=0)
        else:
            endP = domi_dec[torch.randperm(domiN)[:ref_no]]
        ref_no = endP.size(0)
        vector = endP - startP
        Direct = vector / (torch.sum(vector ** 2, dim=1, keepdim=True) ** 0.5)
        nan_index = torch.isnan(Direct)
        Direct[nan_index] = torch.rand(Direct.shape, dtype=torch.double)[nan_index]
        PopDec = []
        for i in range(ref_no):
            lambda1 = (non_dec - startP) @ Direct[i]
            sigma = 1 if len(lambda1) == 1 else torch.std(lambda1) * (1 + ((domiN + 1) / N))
            OffDec = torch.normal(0, sigma, size=(sub_N, 1)) * Direct[i] + startP
            PopDec.append(OffDec)
        OffDec = problem.repair_decision(torch.cat(PopDec, dim=0))
        return problem.estimate_population(PolynomialMutation.do(OffDec, problem))
